from collections import defaultdict
import numpy as np


class Connection:

    def __init__(self, connParams):
        # Paramters
        self.columnCount = connParams['columnCount']
        self.cellsPerColumn = connParams['cellsPerColumn']
        self.initialPermanence = connParams['initialPermanence']
        self.connectedPermanence = connParams['connectedPermanence']
        self.permanenceIncrement = connParams['permanenceIncrement']
        self.permanenceDecrement = connParams['permanenceDecrement']
        self.permanencePunishment = connParams['permanencePunishment']
        self.activationThreshold = connParams['activationThreshold']
        self.minThreshold = connParams['minThreshold']
        self.sampleSize = connParams['sampleSize']
        self.maxSynapsesPerSegment = connParams['maxSynapsesPerSegment']
        # Initializations
        self.miniColumns = [MiniColumn(colId, self.cellsPerColumn) for colId in range(self.columnCount)]
        self.allCells = [cell for miniColumn in self.miniColumns for cell in miniColumn.cells]
        self.allSegments = []
        self.allSynapses = []
        self.nextSegmentId = 0
        self.nextSynapseId = 0
        self.presynToPostsynMap = defaultdict(set)  # from presynapse to postsynapse

    # Low Level Interfaces: Start   ####################################
    @staticmethod
    def filterOverlapByThreshold(overlap, threshold):
        return np.array([s for s in overlap.keys() if overlap[s] >= threshold])

    @staticmethod
    def filterOverlapByInterval(overlap, minValue, maxValue):
        return np.array([s for s in overlap.keys() if minValue <= overlap[s] < maxValue])

    @staticmethod
    def getCellsInColumn(columnId, cellsPerColumn):
        return np.array([co for co in range(columnId * cellsPerColumn, (columnId + 1) * cellsPerColumn)])

    @staticmethod
    def getCellsInColumns(columnIds, cellsPerColumn):
        return np.array([Connection.getCellsInColumn(ci, cellsPerColumn) for ci in columnIds]).flatten()

    def mapCellIdsToColumnIds(self, cellIds):
        return np.array([self.allCells[cellId].miniColumn.columnId for cellId in cellIds])

    def mapSegmentsToCells(self, segmentIds):
        segmentIds = np.array(segmentIds, dtype='int64')
        return np.array([self.allSegments[s].cell.cellId for s in segmentIds],  dtype="uint32")

    def mapCellIdToColumnId(self, cellId):
        return self.allCells[cellId].miniColumn.columnId

    def adjustActiveSynapses(self, learningSegments, presynapseSet, permanenceIncrement, permanenceDecrement):
        for ls in learningSegments:
            self.allSegments[ls].reinforceSynapses(presynapseSet, permanenceIncrement, permanenceDecrement)

    def filterSegmentsByCell(self, segmentIds, cellIds):
        segmentIds = np.array(segmentIds, dtype='int64')
        cellIds = np.array(cellIds, dtype='int64')
        return np.array([so for so in segmentIds if self.allSegments[so].cell.cellId in cellIds])

    def getCellIdBySegmentId(self, segmentId):
        return self.allSegments[segmentId].cell.cellId

    def getSegmentIdsByCellId(self, cellId):
        return np.array([segment.segmentId for segment in self.allCells[cellId].segments])

    def mapSegmentsToSynapseCounts(self, learningSegments):
        return np.array([len(self.allSegments[s].synapses) for s in learningSegments])

    def computeActivity(self, presynapseSet):
        connectedNumOfSegment = defaultdict(int)
        potentialNumOfSegment = defaultdict(int)
        for bi in presynapseSet:
            for synId in self.presynToPostsynMap[bi]:
                synapse = self.allSynapses[synId]
                segment = synapse.segment
                if synapse.permanence >= self.connectedPermanence:
                    connectedNumOfSegment[segment.segmentId] += 1
                if synapse.permanence > 0:
                    potentialNumOfSegment[segment.segmentId] += 1
        return connectedNumOfSegment, potentialNumOfSegment

    def growSynapsesToSample(self, learningSegments, growthCandidates, maxNew, initialPermanece):
        for i, segmentId in enumerate(learningSegments):
            segment = self.allSegments[segmentId]
            presynNotConnectedToSeg = segment.findPresynWithNoConnection(growthCandidates)
            numToSelect = min(maxNew[i], len(presynNotConnectedToSeg))
            if numToSelect > 0:
                np.random.shuffle(presynNotConnectedToSeg)
                presynToGrow = presynNotConnectedToSeg[0:numToSelect]
                self.createSynapses(presynToGrow, segment, initialPermanece)

    def createSegments(self, cellIds):
        cellIds = np.array(cellIds, dtype='int64')
        res = []
        for cellId in cellIds:
            cell = self.allCells[cellId]
            newSegment = Segment(self.nextSegmentId, cell)
            cell.addSegment(newSegment)
            self.allSegments.append(newSegment)
            self.nextSegmentId += 1
            res.append(newSegment.segmentId)
        return np.array(res)

    def createSynapses(self, presynSet, segment, initialPermanane):
        for presyn in presynSet:
            newSyn = Synapse(self.nextSynapseId, segment, presyn, initialPermanane)
            segment.addSynapse(newSyn)
            self.allSynapses.append(newSyn)
            self.presynToPostsynMap[presyn].add(newSyn.synapseId)
            self.nextSynapseId += 1

    # Low Level Interfaces: End     ###################################

    # High Level Interfaces: Start  ###################################
    def createOneSegmentPerCell(self):
        for cell in self.allCells:
            self.createSegments([cell.cellId])

    def getCellsWithFewestSegments(self, columnIds):
        return np.array([self.miniColumns[colummId].getLeastUsedCellIdInColum() for colummId in columnIds])

    def chooseBestSegmentPerColumn(self, matchingCellsInBurstingColumn, allMatchingSegments, potentialOverlaps):
        candidateSegmentIds = self.filterSegmentsByCell(allMatchingSegments, matchingCellsInBurstingColumn)
        maxMatchingSegmentPerCol = defaultdict(int)
        maxPotentialCountPerCol = defaultdict(int)
        for segmentId in potentialOverlaps.keys():
            if segmentId not in candidateSegmentIds:
                continue
            cell = self.allSegments[segmentId].cell
            column = self.allCells[cell.cellId].miniColumn
            count = potentialOverlaps[segmentId]
            if count > maxPotentialCountPerCol[column.columnId]:
                maxPotentialCountPerCol[column.columnId] = count
                maxMatchingSegmentPerCol[column.columnId] = segmentId
        return np.array([key for key in maxMatchingSegmentPerCol.values()])

    def getBestSegmentPerCell(self, candidateSegmentIds, potentialOverlaps):
        maxMatchingSegmentPerCell = defaultdict(int)
        maxPotentialCountPerCell = defaultdict(int)
        for segmentId in potentialOverlaps.keys():
            if segmentId not in candidateSegmentIds:
                continue
            cell = self.allSegments[segmentId].cell
            count = potentialOverlaps[segmentId]
            if count > maxPotentialCountPerCell[cell.cellId]:
                maxPotentialCountPerCell[cell.cellId] = count
                maxMatchingSegmentPerCell[cell.cellId] = segmentId
        return np.array([key for key in maxMatchingSegmentPerCell.values()])

    def getCellsWithNoInhibition(self, partlyDepolarizedCells, fullyDepolarizedCells):
        res = np.array([])
        columnsHasFullyDepoCells = self.mapCellIdsToColumnIds(fullyDepolarizedCells)
        for cellId in partlyDepolarizedCells:
            colId = self.mapCellIdToColumnId(cellId)
            if colId not in columnsHasFullyDepoCells:
                res = np.append(res, cellId)
        return np.array(res)

    def getActiveSegments(self, connectedOverlaps):
        return self.filterOverlapByThreshold(connectedOverlaps, self.activationThreshold)

    def getMatchingSegments(self, potentialOverlaps):
        return self.filterOverlapByThreshold(potentialOverlaps, self.minThreshold)

    def learnOnExisting(self, learningSegments, activeInput, growthCandidates):
        if len(learningSegments) == 0:
            return
        self.adjustActiveSynapses(learningSegments, activeInput, self.permanenceIncrement, self.permanenceDecrement)
        if len(growthCandidates) == 0:
            return
        numNewSynapses = len(growthCandidates)
        if self.sampleSize != -1:
            numNewSynapses = min(self.sampleSize, len(growthCandidates))
        numNewSynapses = [numNewSynapses] * len(learningSegments)
        if self.maxSynapsesPerSegment != -1:
            synapseCounts = self.mapSegmentsToSynapseCounts(learningSegments)
            numSynapsesToReachMax = self.maxSynapsesPerSegment - synapseCounts
            numNewSynapses = np.where(numNewSynapses <= numSynapsesToReachMax, numNewSynapses, numSynapsesToReachMax)
        self.growSynapsesToSample(learningSegments, growthCandidates, numNewSynapses, self.initialPermanence)

    def learnOnNew(self, newSegmentCells, growthCandidates):
        if len(growthCandidates) == 0 or len(newSegmentCells) == 0:
            return
        numNewSynapses = len(growthCandidates)
        if self.sampleSize != -1:
            numNewSynapses = min(numNewSynapses, self.sampleSize)
        if self.maxSynapsesPerSegment != -1:
            numNewSynapses = min(numNewSynapses, self.maxSynapsesPerSegment)
        newSegmentIds = self.createSegments(newSegmentCells)
        maxNew = [numNewSynapses] * len(newSegmentIds)
        self.growSynapsesToSample(newSegmentIds, growthCandidates, maxNew, self.initialPermanence)
        return newSegmentIds

    # def learnInColumnPooler(self, activeCells, presynInput):
    #     for cell in activeCells:
    #         assert len(cell.segments) == 1  # assuming only one segment per cell, for column pooler
    #         theSegment = cell.segments[0]
    #         alreadyConnNonzeroPresyn = theSegment.findAlreadyConnNonzeroPresyn()
    #         synToIncrease = np.array(np.intersect1d(alreadyConnNonzeroPresyn, presynInput))
    #         for sc in synToIncrease:
    #             self.allSynapses[sc].permanence += self.permanenceIncrement
    #             if self.allSynapses[sc].permanence > 1:
    #                 self.allSynapses[sc].permanence = 1
    #         synToDecrease = np.array(np.setdiff1d(alreadyConnNonzeroPresyn, presynInput))
    #         for sd in synToDecrease:
    #             self.allSynapses[sd].permanence -= self.permanenceDecrement
    #             if self.allSynapses[sd].permanence < 0:
    #                 self.allSynapses[sd].permanence = 0
    #         synToGrowCandiate = np.array(np.setdiff1d(presynInput, alreadyConnNonzeroPresyn))
    #         numToGrow = min(self.sampleSize, len(synToGrowCandiate))
    #         np.random.shuffle(synToGrowCandiate)
    #         synToGrowActual = synToGrowCandiate[0:numToGrow]
    #         self.createSynapses(synToGrowActual, theSegment, self.initialPermanence)
    # High Level Interfaces: End    ####################################


class MiniColumn:

    def __init__(self, columnId, cellsPerColumn):
        self.columnId = columnId
        self.cellsPerCol = cellsPerColumn
        self.cells = [Cell(cellId, self) for cellId in Connection.getCellsInColumn(columnId, cellsPerColumn)]
        # temporal states, cleared and recalculated in every step
        self.matchingSegmentIds = np.array([])
        self.activeSegmentIds = np.array([])
        self.hasPrediction = False
        self.hasPotential = False
        self.maxMatchingSegment = None
        self.predictedCellIds = np.array([])

    def getLeastUsedCellIdInColum(self):
        leastSegCellIds = []
        minNumSegments = float("inf")
        for cell in self.cells:
            numOfSegsOfCell = len(cell.segments)
            if numOfSegsOfCell < minNumSegments:
                minNumSegments = numOfSegsOfCell
                leastSegCellIds = []
            if numOfSegsOfCell == minNumSegments:
                leastSegCellIds.append(cell.cellId)
        i = np.random.randint(0, len(leastSegCellIds))
        return leastSegCellIds[i]

    def clear(self):
        self.matchingSegmentIds = np.array([])
        self.activeSegmentIds = np.array([])
        self.hasPrediction = False
        self.hasPotential = False
        self.maxMatchingSegment = None
        self.predictedCellIds = np.array([])


class Cell:

    def __init__(self, cellId, miniColumn):
        self.cellId = cellId
        self.miniColumn = miniColumn
        self.segments = []

    def addSegment(self, segment):
        self.segments.append(segment)


class Segment:

    def __init__(self, segmentId, cell):
        self.segmentId = segmentId
        self.cell = cell
        self.synapses = []

    def addSynapse(self, synapse):
        self.synapses.append(synapse)

    def reinforceSynapses(self, presynapseSet, permIncr, permDecr):
        for syn in self.synapses:
            if syn.preSynCell in presynapseSet:
                syn.permanence += permIncr  # reward synapses with active presynaptic cells
            else:
                syn.permanence -= permDecr
            syn.permanence = max(0.0, min(1.0, syn.permanence))
            syn.permanence = syn.permanence

    def findPresynWithNoConnection(self, presynapseSet):
        alreadyConnNonzeroPresynSet = self.findAlreadyConnNonzeroPresyn()
        return np.array([presyn for presyn in presynapseSet if presyn not in alreadyConnNonzeroPresynSet])

    def findAlreadyConnNonzeroPresyn(self):
        return set([s.preSynCell for s in self.synapses if s.permanence > 0])


class Synapse:

    def __init__(self, synapseId, segment, preSynCell, permanence):
        self.synapseId = synapseId
        self.segment = segment
        self.preSynCell = preSynCell
        self.permanence = permanence
